(in-package :context-indexing)

(defvar *standard-optimize-settings* '(optimize (speed 3) 
                                                (safety 0) 
                                                (compilation-speed 0) 
                                                (space 0) 
                                                (debug 0)))

(deftype particle () '(simple-array single-float))

(defmacro signature (function arguments result)
  `(declaim (ftype (function ,arguments ,result) ,function)))

(signature pull (particle particle single-float) particle)
(signature center ((simple-array particle) 
                   &optional (simple-array (single-float 0.0))) 
                  particle)
(signature euclidian-distance (particle particle) (single-float 0.0))
(signature mean (particle) single-float)
(signature central-moment (fixnum particle) single-float)
(signature kurtosis (particle) (single-float 0.0))

(declaim (inline l2-norm 
                 nnormalize
                 nscale
                 central-moment
                 mean
                 make-particle))

(defun starts-with (list head)
  (and (consp list) (eql (first list) head)))

(defun assoc-values (keys alist &key (test #'eql))
  (mapcar (op cdr (assoc _ alist :test test)) keys))
  
(define-compiler-macro map-into1 (function result-sequence &rest sequences)
  (let ((result (gensym)))
    `(let ((,result ,result-sequence))
       (map-into ,result ,function ,result ,@sequences)
       ,result)))
      
(define-compiler-macro foldl (function initial-value &rest sequences)
  (let* ((fold (gensym))
         (body `(funcall ,function ,fold ,@(mapcar (constantly '_) sequences)))
         (result-type (if (starts-with initial-value 'the)
                          (second initial-value)
                          t)))
    `(let ((,fold ,initial-value))
       ,@(unless (eq result-type t) `((declare (type ,result-type ,fold))))
       (map nil ,(if (and (starts-with function 'function) 
                          (get (second function) 'mutablep))
                     `(op ,@body)
                     `(op setf ,fold ,body))
            ,@sequences)
       ,fold)))

(define-compiler-macro sum (function &rest sequences)
  (if sequences
      `(foldl (op + _ (funcall ,function ,@(mapcar (constantly '_) sequences)))
              (the ,(if (and (starts-with function 'op)
                             (eq (second function) 'expt)
                             (integerp (fourth function))
                             (evenp (fourth function)))
                        '(single-float 0.0) 
                        'single-float)
                   0.0)
              ,@sequences)
      `(foldl #'+ (the single-float 0.0) ,function)))

(defun unfold-restarg (rule form)
  (destructuring-bind (names function sequences) rule
    (aif (second (member '&rest names))
         (let ((unfold (loop :repeat (- (length form) (length names) -1) 
                             :collect (gensym))))
           (list (append (butlast names 2) unfold)
                 (subst unfold it function)
                 (cond ((consp sequences) (subst unfold it sequences)) 
                       ((eq sequences it) unfold)
                       (t sequences))))
         rule)))
                 
(defun unify (rule form)
  (destructuring-bind (names function sequences) (unfold-restarg rule form)
    (let ((binds (pairlis names (rest form))))
      (values (sublis binds function)
              (assoc-values sequences binds) 
              (assoc-values (set-difference names sequences) binds)))))

(defun map-inline-rule (form &key recursivep)
  (flet ((nth-rule (n) (op nth-value n (map-inline-rule _ :recursivep t))))
    (aif (and (listp form) 
              (not (and recursivep (get (first form) 'mutablep)))
              (get (first form) 'map-inline-rule))
         (multiple-value-bind (function sequences invariants) (unify it form)
           (values (mapcar (nth-rule 0) function)
                   (remove-duplicates (mapcan (nth-rule 1) sequences)
                                      :from-end t :test #'equal)
                   (remove-duplicates (append (mapcan (nth-rule 2) form)
                                              invariants) 
                                      :from-end t :test #'equal)))
         (values (if (listp form) (mapcar (nth-rule 0) form) form) 
                 (list form) 
                 nil))))

(defun make-fresh-bindings (forms)
  (mapcar (op* cons _ (gensym)) forms))

(defun fold-funcall (form)
  (if (starts-with form 'funcall) (cons (cadadr form) (cddr form)) form))
 
(defun fold-mutable-calls (form)
  (labels ((fold (form)
             (if (and (listp form) (get (first form) 'map-inline-rule))
                 (destructuring-bind (head first . rest) form
                   (list* (get head 'immutable-dual head) (fold first) rest))
                 form)))
    (list* (first form) (fold (second form)) (cddr form))))
      
(defun pervade-immutability (form)
  (labels ((mutable-call-tree-p (form)
             (or (atom form) 
                 (not (get (first form) 'map-inline-rule))
                 (and (get (first form) 'mutablep)
                      (mutable-call-tree-p (second form))))))
    (if (or (not (get (first form) 'mutablep)) (mutable-call-tree-p form))
        form
        (cons (get (first form) 'immutable-dual) (rest form)))))
 
(defun inline-map (form env)
  (declare (ignore env))
  (let ((form (fold-mutable-calls (pervade-immutability (fold-funcall form)))))
    (multiple-value-bind (function sequences invariants) (map-inline-rule form)
      (let ((args (make-fresh-bindings sequences))
            (invariants (make-fresh-bindings invariants)))
        `(,@(if (get (first form) 'mutablep) '(map-into1) '(map 'vector))
             (let ,(loop :for (form . name) :in invariants 
                         :collect `(,name ,form))
               (lambda ,(mapcar #'cdr args) 
                 ,(sublis (append invariants args) function :test #'equal)))
             ,@sequences)))))
 
(defmacro define-map-inline ((head . names) function sequences &key mutablep)
  `(progn
     (setf (get ',head 'map-inline-rule) (list ',names ',function ',sequences))
     (setf (compiler-macro-function ',head) #'inline-map)
     (setf (get ',head 'mutablep) ,mutablep)
     ,@(when mutablep
         (let ((dual (gensym))) 
           `((define-map-inline (,dual ,@names) ,function ,sequences)
             (setf (get ',head 'immutable-dual) ',dual))))))

(define-map-inline (vector- &rest vectors) (- . vectors) vectors)
(define-map-inline (scale vector scalar) (* vector scalar) (vector))
(define-map-inline (nvector+ &rest vectors) (+ . vectors) vectors :mutablep t)
(define-map-inline (nvector+* a b scalar) (+ a (* b scalar)) (a b) :mutablep t)
